import os
import subprocess
import torch
print(torch.__version__)  # 应该输出版本号，如 2.0.0
print(torch.cuda.is_available())  # 检查 GPU 是否可用（True 表示可用）

# 设置环境变量
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# 你的训练命令
train_command = 'yolo train model=yolov8n.pt data="D:/yolov8_2/ultralytics-main/custom_dataset.yaml" epochs=10 verbose=True batch=15 device=0'

# 执行训练
subprocess.call(train_command, shell=True)
